import numpy as np
from numba import njit
import pybroker
from pybroker import Strategy
from pybroker.ext.data import AKShare

import talib  
pybroker.enable_data_source_cache('AKShare')
 
def ma(bar_data,short_window,long_window):
    
    # 计算短期（例如5日）和长期（例如20日）均线  
    df['SMA_short'] = talib.SMA(df['close'], timeperiod=short_window)  
    df['SMA_long'] = talib.SMA(df['close'], timeperiod=long_window)  
  
    # 初始化金叉和死叉的列  
    df['GoldenCross'] = 0  
    df['DeathCross'] = 0  
  
    # 遍历DataFrame，判断金叉和死叉  
    for i in range(1, len(df)):  
        if df['SMA_short'].iloc[i] > df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] < df['SMA_long'].iloc[i-1]:  
            df.at[df.index[i], 'GoldenCross'] = 1  
        if df['SMA_short'].iloc[i] < df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] > df['SMA_long'].iloc[i-1]:  
            df.at[df.index[i], 'DeathCross'] = -1  
            
    return

aKShare = AKShare()
df = aKShare.query('000001', '4/1/2021', '4/1/2022')

ma(df.close, 5, 20)  
# 查看结果  
print(df.tail(50))